#include "../include/list.h"
#include "../include/hash.h"
#include "../include/zmsg.h"

#include "../include/zbusapi.h"

struct _zbuscli_cfg_t{
	void* 	ctx;
	int 	verbose;
	char*	broker;
	int		timeout;
	int		retries;
	int		io_threads;
};

struct _zbuscli_t{
	void* 	ctx;
	int		own_ctx;

	void* 	socket;
	int 	verbose;
	char*	broker;
	int		timeout;
	int		retries;
};

struct _zbuswrk_cfg_t{
	//the followings for worker_pool
	void* 	ctx;
	int		io_threads;
	list_t*	brokers;

	char*	broker;
	char*	service;
	char*	mode;
	int 	verbose;

	char*	register_token;
	char*	access_token;

	int		heartbeat;
	int		reconnect;
	int		timeout; 
};

struct _zbuswrk_t{
	char*	broker;
	char*	service;
	char*	mode;

	char*	register_token;
	char*	access_token;

	int64_t heartbeat_at;

	int 	liveness;
	int		heartbeat;
	int		reconnect;
	int		timeout;

	void* 	ctx;
	void* 	socket;
	int 	verbose;
	int		own_ctx;

	zframe_t* reply_to;
};

static void
cli_connect_to_broker(zbuscli_t* self){
	assert(self);
	assert(self->ctx);
	int rc;
	if(self->socket){
		rc = zmq_close(self->socket);
		assert( rc == 0);
	}
	self->socket = zmq_socket(self->ctx, ZMQ_REQ);
	assert(self->socket);
	int linger = 0;
	rc = zmq_setsockopt(self->socket, ZMQ_LINGER, &linger, sizeof(linger));
	assert( rc == 0);
	rc = zmq_connect(self->socket, self->broker);
	assert( rc == 0);

	if(self->verbose){
		zlog("I: connecting to broker at %s... \n", self->broker);
	}
}

zbuscli_t*
zbuscli_new(zbuscli_cfg_t* cfg){
	assert(cfg);
	assert(cfg->io_threads>0);

	zbuscli_t* self = (zbuscli_t*)zmalloc(sizeof(zbuscli_t));
	memset(self, 0, sizeof(zbuscli_t));
	if(cfg->ctx){
		self->ctx = cfg->ctx;
		self->own_ctx = 0;
	} else {
		self->ctx = zctx_new(cfg->io_threads);
		self->own_ctx = 1;
	}
	self->retries = cfg->retries;
	self->verbose = cfg->verbose;
	self->timeout = cfg->timeout;
	self->broker = zstrdup(cfg->broker);
	cli_connect_to_broker(self);

	return self;
}

void
zbuscli_destroy(zbuscli_t** self_p){
	assert( self_p );
	zbuscli_t* self = *self_p;
	if(self){
		if(self->socket){
			int rc = zmq_close(self->socket);
			assert(rc == 0);
		}
		if(self->broker){
			zfree(self->broker);
		}
		if(self->own_ctx && self->ctx){
			zctx_destroy(&self->ctx);
		}
		zfree(self);
		*self_p = NULL;
	}
}

zbuscli_cfg_t*
zbuscli_cfg_new(){
	zbuscli_cfg_t* self = (zbuscli_cfg_t*)zmalloc(sizeof(zbuscli_cfg_t));
	assert(self);
	memset(self, 0, sizeof(zbuscli_cfg_t));
	self->broker = zstrdup("tcp://localhost:15555"); //default;
	self->io_threads = 1;
	self->retries = 1;
	self->timeout = 2500;//ms
	self->verbose = 0;

	return self;
}

void
zbuscli_cfg_destroy(zbuscli_cfg_t** self_p){
	assert(self_p);
	zbuscli_cfg_t* self = *self_p;
	if(self){
		if(self->broker)
			zfree(self->broker);
		zfree(self);
		*self_p = NULL;
	}
}
void
zbuscli_cfg_set_ctx(zbuscli_cfg_t* self, void* ctx){
	assert(self);
	self->ctx = ctx;
}

void
zbuscli_cfg_set_verbose(zbuscli_cfg_t* self, int verbose){
	assert(self);
	self->verbose = verbose;
}

void
zbuscli_cfg_set_timeout(zbuscli_cfg_t* self, int timeout){
	assert(self);
	self->timeout = timeout;
}

void
zbuscli_cfg_set_retries(zbuscli_cfg_t* self, int retries){
	assert(self);
	self->retries = retries;
}
void
zbuscli_cfg_set_iothreads(zbuscli_cfg_t* self, int io_threads){
	assert(self);
	self->io_threads = io_threads;
}

void
zbuscli_cfg_set_broker(zbuscli_cfg_t* self, char* host, int port){
	assert(self);
	assert(host);
	if(self->broker)
		zfree(self->broker);
	char broker[256];
	sprintf(broker, "tcp://%s:%d", host, port);
	self->broker = zstrdup(broker);
}


zmsg_t*
zbuscli_request(zbuscli_t* self, char* service, char* token, zmsg_t* request){
	assert(self);
	assert(request);
	zmsg_push_front(request, zframe_newstr(token));
	zmsg_push_front(request, zframe_newstr(service));
	zmsg_push_front(request, zframe_newstr(MDPC_CLIENT));
	if(self->verbose){
		zlog("I: send request to '%s' service\n", service);
		zmsg_log(request);
	}
	int rc = zmsg_send(&request, self->socket);
	if(rc != 0){
		return NULL;
	}

	int retries = self->retries;
	while( retries>0 ) {
		zmq_pollitem_t items [] = { { self->socket,  0, ZMQ_POLLIN, 0 } };
		int rc = zmq_poll (items, 1, self->timeout * ZMQ_POLL_MSEC);
		if(rc == -1){//interrupted
			return NULL;
		}
		if (items [0].revents & ZMQ_POLLIN) {
			zmsg_t* msg = zmsg_recv(self->socket);
			if(!msg) break; //interrupted
			if(self->verbose){
				zmsg_log(msg);
			}
			assert( zmsg_frame_size(msg)>=3 );
			zframe_t* header = zmsg_pop_front(msg);
			zframe_destroy(&header);

			zframe_t* service = zmsg_pop_front(msg);
			zframe_destroy(&service);

			return msg;
		} else {
			cli_connect_to_broker(self);
			if(--retries > 0){
				if(self->verbose){
					zlog("I: no reply, reconnecting... \n");
				} 
			} else {
				if(self->verbose){
					zlog("I: no reply, abandoning\n");
				}
			}
		}
	}
	return NULL;
}
/////////////////////////////////////////////////////////////////////////
static void
destroy_str(void** ptr){
	if(ptr && *ptr){
		zfree(*ptr);
		*ptr = NULL;
	}
}

zbuswrk_cfg_t*
zbuswrk_cfg_new(){
	zbuswrk_cfg_t* self = (zbuswrk_cfg_t*)zmalloc(sizeof(zbuswrk_cfg_t));
	assert(self);
	memset(self, 0, sizeof(zbuswrk_cfg_t));

	//set default values
	self->broker = zstrdup("tcp://localhost:15555");
	self->mode = zstrdup(MODE_REQREP);
	self->heartbeat = 2500;
	self->reconnect = 2500;
	self->timeout	= 2500;
	self->verbose = 0;
	self->access_token = zstrdup("");
	self->register_token = zstrdup("");

	//for pool
	self->io_threads = 1;
	self->brokers = list_new();
	list_set_destroy(self->brokers, destroy_str);

	return self;
}
//pool part not copy
zbuswrk_cfg_t*
zbuswrk_cfg_dup(zbuswrk_cfg_t* from){
	zbuswrk_cfg_t* self = (zbuswrk_cfg_t*)zmalloc(sizeof(zbuswrk_cfg_t));
	assert(self);
	memset(self, 0, sizeof(zbuswrk_cfg_t));
	//copy values
	self->service = zstrdup(from->service); 
	self->ctx 	  = from->ctx;

	self->broker = zstrdup(from->broker);
	self->mode = zstrdup(from->mode);
	self->heartbeat = from->heartbeat;
	self->reconnect = from->reconnect;
	self->timeout	= from->timeout;
	self->verbose = from->verbose;
	self->access_token = zstrdup(from->access_token);
	self->register_token = zstrdup(from->register_token);

	//pool part not copy
	return self;
}

void
zbuswrk_cfg_destroy(zbuswrk_cfg_t** self_p){
	assert(self_p);
	zbuswrk_cfg_t* self = *self_p;
	if(self){
		if(self->broker)
			zfree(self->broker);
		if(self->mode)
			zfree(self->mode);
		if(self->register_token)
			zfree(self->register_token);
		if(self->access_token)
			zfree(self->access_token);

		if(self->service)
			zfree(self->service);

		if(self->brokers)
			list_destroy(&self->brokers);

		zfree(self);
		*self_p = NULL;
	}
}

void
zbuswrk_cfg_set_ctx(zbuswrk_cfg_t* self, void* ctx){
	assert(self);
	self->ctx = ctx;
}
void
zbuswrk_cfg_set_verbose(zbuswrk_cfg_t* self, int verbose){
	assert(self);
	self->verbose = verbose;
}

void
zbuswrk_cfg_set_reconnect(zbuswrk_cfg_t* self, int reconnect){
	assert(self);
	self->reconnect = reconnect;
}

void
zbuswrk_cfg_set_timeout(zbuswrk_cfg_t* self, int timeout){
	assert(self);
	self->timeout = timeout;
}

void
zbuswrk_cfg_set_heartbeat(zbuswrk_cfg_t* self, int heartbeat){
	assert(self);
	self->heartbeat = heartbeat;
}

void
zbuswrk_cfg_set_service(zbuswrk_cfg_t* self, char* servcie){
	assert(self);
	if(self->service)
		zfree(self->service);
	self->service = zstrdup(servcie);
}
void
zbuswrk_cfg_set_regtoken(zbuswrk_cfg_t* self, char* register_token){
	assert(self);
	if(self->register_token)
		zfree(self->register_token);
	self->register_token = zstrdup(register_token);
}
void
zbuswrk_cfg_set_acctoken(zbuswrk_cfg_t* self, char* access_token){
	assert(self);
	if(self->access_token)
		zfree(self->access_token);
	self->access_token = zstrdup(access_token);
}

void
zbuswrk_cfg_set_mode(zbuswrk_cfg_t* self, char* mode){
	assert(self);
	if(self->mode)
		zfree(self->mode);
	self->mode = zstrdup(mode);
}

void
zbuswrk_cfg_set_broker(zbuswrk_cfg_t* self, char* broker){
	assert(self);
	if(self->broker){
		zfree(self->broker);
	}
	self->broker = zstrdup(broker);
}
void
zbuswrk_cfg_set_broker2(zbuswrk_cfg_t* self, char* host, int port){
	assert(self);
	if(self->broker)
		zfree(self->broker);
	char broker[256];
	sprintf(broker, "tcp://%s:%d", host, port);
	self->broker = zstrdup(broker);
}

void
zbuswrk_cfg_add_broker(zbuswrk_cfg_t* self, char* host, int port){
	assert(self);
	char broker[256];
	sprintf(broker, "tcp://%s:%d",host, port);
	list_push_back(self->brokers, zstrdup(broker));
}
 

static void
wrk_send_to_broker(zbuswrk_t* self, char* command, char* option, zmsg_t* msg){
	if(msg == NULL){
		msg = zmsg_new();
	}
	if(option){
		zmsg_push_front_str(msg, option);
	}
	zmsg_push_front_str(msg, command);
	zmsg_push_front_str(msg, MDPW_WORKER);
	zmsg_push_front(msg, zframe_newstr(""));
	zmsg_send(&msg, self->socket);
}

static void
wrk_connect_to_broker(zbuswrk_t* self){
	assert(self);
	int rc;
	if(self->socket){
		rc = zmq_close(self->socket);
		assert(rc == 0);
	}
	self->socket = zmq_socket(self->ctx, ZMQ_DEALER);
	assert(self->socket);
	int linger = 0;
	rc = zmq_setsockopt(self->socket, ZMQ_LINGER, &linger, sizeof(linger));
	assert(rc == 0);
	rc = zmq_connect(self->socket, self->broker);
	assert(rc == 0);

	if(self->verbose){
		zlog("I: connect to broker at %s ....\n", self->broker);
	}

	zmsg_t* msg = zmsg_new();
	zmsg_push_back_str(msg, self->register_token);
	zmsg_push_back_str(msg, self->access_token);
	zmsg_push_back_str(msg, self->mode);

	wrk_send_to_broker(self, MDPW_READY, self->service, msg);

	self->liveness = HEARTBEAT_LIVENESS;
	self->heartbeat_at = zclock_time()+self->heartbeat;
}


zbuswrk_t*
zbuswrk_new(zbuswrk_cfg_t* cfg){
	assert(cfg);
	assert(cfg->service);
	zbuswrk_t* self = (zbuswrk_t*)zmalloc(sizeof(zbuswrk_t));
	memset(self, 0, sizeof(zbuswrk_t));

	if(cfg->ctx){
		self->ctx = cfg->ctx;
		self->own_ctx = 0;
	} else {
		self->ctx = zctx_new(1);
		self->own_ctx = 1;
	}
	self->verbose = cfg->verbose;
	self->timeout  = cfg->timeout;
	self->heartbeat = cfg->heartbeat;
	self->reconnect = cfg->reconnect;

	self->broker = zstrdup(cfg->broker);
	self->service = zstrdup(cfg->service);
	self->mode    = zstrdup(cfg->mode);
	self->register_token = zstrdup(cfg->register_token);
	self->access_token = zstrdup(cfg->access_token);

	wrk_connect_to_broker(self);

	return self;
}

void
zbuswrk_destroy(zbuswrk_t** self_p){
	assert(self_p);
	zbuswrk_t* self = *self_p;
	int rc;
	if(self){
		if(self->broker){
			zfree(self->broker);
		}
		if(self->service){
			zfree(self->service);
		}
		if(self->mode){
			zfree(self->mode);
		}
		if(self->register_token){
			zfree(self->register_token);
		}
		if(self->access_token){
			zfree(self->access_token);
		}
		if(self->reply_to){
			zframe_destroy(&self->reply_to);
		}
		if(self->socket){
			rc = zmq_close(self->socket);
			assert(rc == 0);
		}
		if(self->own_ctx && self->ctx){
			zctx_destroy(&self->ctx);
		}
		zfree(self);
		*self_p = NULL;
	}
}


zmsg_t*
zbuswrk_recv(zbuswrk_t* self, zmsg_t* reply){
	if(reply){
		assert(self->reply_to);
		zmsg_wrap(reply, self->reply_to);
		wrk_send_to_broker(self, MDPW_REPLY, NULL, reply);
	}

	while(1) {
		zmq_pollitem_t items [] = { { self->socket,  0, ZMQ_POLLIN, 0 } };
		int rc = zmq_poll (items, 1, self->timeout * ZMQ_POLL_MSEC);
		if(rc == -1){//interrupted
			return NULL;
		}
		if (items [0].revents & ZMQ_POLLIN) {
			zmsg_t* msg = zmsg_recv(self->socket);
			if(!msg) break; //interrupted
			
			self->liveness = HEARTBEAT_LIVENESS;

			assert( zmsg_frame_size(msg)>=3 );

			zframe_t* empty = zmsg_pop_front(msg);
			assert(zframe_streq(empty, ""));
			zframe_destroy(&empty);

			zframe_t* header = zmsg_pop_front(msg);
			assert(zframe_streq(header, MDPW_WORKER));
			zframe_destroy(&header);

			zframe_t* cmd = zmsg_pop_front(msg);
			if(zframe_streq(cmd, MDPW_REQUEST)){
				if(self->verbose){
					zmsg_log(msg);
				}
				self->reply_to = zmsg_unwrap(msg);
				zframe_destroy(&cmd);

				return msg;
			} else if(zframe_streq(cmd, MDPW_HEARTBEAT)){
				// Do nothing for heartbeats
			} else if(zframe_streq(cmd, MDPW_DISCONNECT)){
				if(self->verbose){
					zlog("I: disconnected by peer broker\n");
					zmsg_log(msg);
				}
				zframe_destroy(&cmd);
				zmsg_destroy(&msg);
				break;
			} else if(zframe_streq(cmd, MDPW_SYNC)){
				if(self->verbose){
					zlog("I: synchronise reset by peer broker\n");
					zmsg_log(msg);
				}
				wrk_connect_to_broker(self);
			} else {
				zlog("E: invalid request message\n");
				zmsg_log(msg);
			}

			zframe_destroy(&cmd);
			zmsg_destroy(&msg);

		} else if( --self->liveness == 0){
			if(self->verbose){
				zlog("W: disconnected from broker(liveness down) - retrying...\n");
			}
			zclock_sleep(self->reconnect);
			wrk_connect_to_broker(self);
		}

		if( zclock_time() > self->heartbeat_at){
			wrk_send_to_broker(self, MDPW_HEARTBEAT, NULL, NULL);
			self->heartbeat_at = zclock_time() + self->heartbeat;
		}
	}
	return NULL;
}


/////////////////////////////////////////worker pool///////////////////////////////////
typedef struct _cfg_msgcb_t{
	zbuswrk_cfg_t* cfg;
	msg_cb cb;
} cfg_msgcb_t; 

typedef struct _cfg_wrkcb_t{
	zbuswrk_cfg_t* cfg;
	worker_cb cb;
} cfg_wrkcb_t;

void
handle_destroy(void** ptr){
	if(ptr && *ptr) zfree(*ptr);
}
 
static void*
s_thread_msg_cb(void* args){
	cfg_msgcb_t* s = (cfg_msgcb_t*)args;
	zbuswrk_cfg_t* cfg = s->cfg;
	msg_cb cb = s->cb;

	zbuswrk_t* worker = zbuswrk_new(cfg);
	zmsg_t* reply = NULL;
	while(1){
		zmsg_t* request = zbuswrk_recv(worker, reply);
		if(!request) break;

		reply = cb(request);

		if(!reply){ //echo empty back,anyway
			reply = zmsg_new();
		} 
	}
	zbuswrk_destroy(&worker);
	zbuswrk_cfg_destroy(&cfg);
	zfree(s);

	return NULL;
}
 

void
zbuswrk_pool_msg_cb(zbuswrk_cfg_t* cfg, int wrk_threads, msg_cb cb){
	assert(cfg);
	assert(cfg->brokers);
	assert(list_size(cfg->brokers));
	assert(cfg->service);
	assert(cfg->io_threads>0);
	int own_ctx = 0;
	if(!cfg->ctx){
		own_ctx = 1;
		cfg->ctx  = zctx_new(cfg->io_threads);
		assert(cfg->ctx );
	}

	list_t* all_threads = list_new();
	list_set_destroy(all_threads, handle_destroy);
	list_node_t* node = list_head(cfg->brokers);
	while(node){
		char* broker = (char*)list_value(node);
		//new thread has to destroy this new config
		int i;
		for(i=0; i<wrk_threads; i++){
			cfg_msgcb_t* s = (cfg_msgcb_t*)zmalloc(sizeof(cfg_msgcb_t));
			zbuswrk_cfg_t* wrk_cfg = zbuswrk_cfg_dup(cfg);
			zbuswrk_cfg_set_broker(wrk_cfg, broker);
			s->cfg = wrk_cfg;
			s->cb = cb; 

			zthread_t* thread = (zthread_t*)zmalloc(sizeof(zthread_t));
			*thread = zthread_new(s_thread_msg_cb, s); 
			list_push_back(all_threads, thread);

		}

		node = list_next(node);
	}

	node = list_head(all_threads);
	while(node){ 
		zthread_t* thread = (zthread_t*)list_value(node);
		zthread_join(*thread);	 
		node = list_next(node);
	}

	list_destroy(&all_threads);

	if(own_ctx)
		zctx_destroy(&cfg->ctx);
}
 



static void*
s_thread_worker_cb(void* args){
	cfg_wrkcb_t* s = (cfg_wrkcb_t*)args;
	zbuswrk_cfg_t* cfg = s->cfg;
	worker_cb cb = s->cb;

	zbuswrk_t* worker = zbuswrk_new(cfg);
	
	cb(worker);

	zbuswrk_destroy(&worker);
	zbuswrk_cfg_destroy(&cfg);
	zfree(s);

	return NULL;
}
  
void
zbuswrk_pool_worker_cb(zbuswrk_cfg_t* cfg, int wrk_threads, worker_cb cb){
	assert(cfg);
	assert(cfg->brokers);
	assert(cfg->service);
	assert(cfg->io_threads>0);
	int own_ctx = 0;
	if(!cfg->ctx){
		own_ctx = 1;
		cfg->ctx  = zctx_new(cfg->io_threads);
		assert(cfg->ctx );
	}

	list_t* all_threads = list_new();
	list_set_destroy(all_threads, handle_destroy);
	list_node_t* node = list_head(cfg->brokers);
	while(node){
		char* broker = (char*)list_value(node);
		//new thread has to destroy this new config
		int i;
		for(i=0; i<wrk_threads; i++){
			cfg_wrkcb_t* s = (cfg_wrkcb_t*)zmalloc(sizeof(cfg_wrkcb_t));
			zbuswrk_cfg_t* wrk_cfg = zbuswrk_cfg_dup(cfg);
			zbuswrk_cfg_set_broker(wrk_cfg, broker);
			s->cfg = wrk_cfg;
			s->cb = cb;

			zthread_t* thread = (zthread_t*)zmalloc(sizeof(zthread_t));
			*thread = zthread_new(s_thread_worker_cb, s); 
			list_push_back(all_threads, thread);
		}

		node = list_next(node);
	}

	node = list_head(all_threads);
	while(node){ 
		zthread_t* thread = (zthread_t*)list_value(node);
		zthread_join(*thread);	 
		node = list_next(node);
	}

	list_destroy(&all_threads);

	if(own_ctx)
		zctx_destroy(&cfg->ctx);
}
